# _*_ coding: utf-8 _*_
# .@FileName:test
# .@Data....:2025-02-23 : 21 : 12
# .@Aurhor..:LiuJingYu
"""
launch:
from __future__ import unicode_literals, print_function
import maya.cmds as mc
"""
import pymel.core as pm
import maya.OpenMaya as om
import math
import logging
from typing import List, Tuple, Optional

# 配置日志记录
# 设置日志记录的基本配置
# level=logging.INFO 表示只记录 INFO 级别及以上的日志信息
# format 定义了日志信息的格式，包含时间、日志名称、日志级别和具体消息
# filename 指定日志文件的名称，这里是 "dis_weight.log"
# filemode="w" 表示以写入模式打开日志文件，每次运行程序会覆盖原有内容
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    filename="dis_weight.log",
    filemode="w"
)
# 创建一个日志记录器，使用当前模块的名称
logger = logging.getLogger(__name__)


def get_distance(start: Tuple[float, float, float],
                 end: Tuple[float, float, float]) -> float:
    """计算两个三维坐标点的欧氏距离（优化版）

    Args:
        start: 起始点坐标 (x, y, z)
        end: 结束点坐标 (x, y, z)

    Returns:
        两点之间的欧氏距离
    """
    # 使用 zip 函数同时遍历 start 和 end 中的对应元素
    # 计算每个对应元素差值的平方和，然后取平方根得到欧氏距离
    return math.sqrt(sum((s - e) ** 2 for s, e in zip(start, end)))


def dis_weight(mesh: pm.nodetypes.Transform,
               def_mesh: pm.nodetypes.Transform,
               joint: pm.nodetypes.Joint) -> int:
    """基于参考网格的蒙皮权重分配工具

    通过计算参考网格与目标网格顶点之间的距离比例，自动分配关节权重

    Args:
        mesh: 目标网格物体（需要分配权重的网格）
        def_mesh: 参考网格物体（提供位置参考的网格）
        joint: 目标关节

    Returns:
        0 表示成功，非零表示错误代码

    Raises:
        RuntimeError: 当输入参数不符合要求时抛出
    """
    try:
        # 参数验证
        # 检查输入的 mesh、def_mesh 是否为 Transform 类型，joint 是否为 Joint 类型
        if not all([isinstance(mesh, pm.nodetypes.Transform),
                    isinstance(def_mesh, pm.nodetypes.Transform),
                    isinstance(joint, pm.nodetypes.Joint)]):
            # 若参数类型错误，记录错误日志
            logger.error("输入参数类型错误")
            # 在 Maya 中显示错误信息
            om.MGlobal.displayError("参数必须是Transform和Joint类型")
            return 1

        # 获取蒙皮簇
        skin_clusters = []
        # 遍历目标网格的所有形状节点
        for shape in mesh.getShapes():
            # 获取每个形状节点历史中的蒙皮簇节点，并添加到 skin_clusters 列表中
            skin_clusters.extend(shape.history(type="skinCluster"))

        if not skin_clusters:
            # 若未找到蒙皮簇，记录错误日志
            logger.error("未找到蒙皮簇")
            # 在 Maya 中显示错误信息
            om.MGlobal.displayError("未找到蒙皮簇")
            return 2

        # 取第一个蒙皮簇作为要使用的蒙皮簇
        skin = skin_clusters[0]
        if len(skin_clusters) > 1:
            # 若发现多个蒙皮簇，记录警告日志
            logger.warning("发现多个蒙皮簇，使用第一个")
            # 在 Maya 中显示警告信息
            om.MGlobal.displayWarning("发现多个蒙皮簇，使用第一个")

        # 验证关节是否为影响物体
        # 获取蒙皮簇的所有影响关节
        influences = skin.influenceObjects()
        if joint not in influences:
            # 若目标关节不在影响关节列表中，记录错误日志
            logger.error(f"关节 {joint} 不是蒙皮影响物体")
            # 在 Maya 中显示错误信息
            om.MGlobal.displayError(f"关节 {joint} 不是蒙皮影响物体")
            return 3

        # 预计算顶点位置（世界空间）
        logger.info("正在获取顶点位置数据...")
        # 获取参考网格所有顶点在世界空间的位置，存储在 def_points 列表中
        def_points = [pm.xform(vtx, q=True, t=True, ws=True)
                      for vtx in def_mesh.vtx]
        # 使用生成器表达式获取目标网格所有顶点在世界空间的位置，节省内存
        mesh_points = (pm.xform(vtx, q=True, t=True, ws=True)
                       for vtx in mesh.vtx)

        # 计算距离列表并获取最大值
        logger.info("正在计算顶点距离...")
        # 计算参考网格和目标网格对应顶点之间的距离，存储在 distances 列表中
        distances = [get_distance(d, m) for d, m in zip(def_points, mesh_points)]

        if not distances:
            # 若距离列表为空，说明顶点数据不匹配，记录错误日志
            logger.error("顶点数据不匹配")
            # 在 Maya 中显示错误信息
            om.MGlobal.displayError("网格顶点数量不一致")
            return 4

        # 获取最大距离
        max_distance = max(distances)
        if max_distance <= 0:
            # 若最大距离无效，记录错误日志
            logger.error("最大距离无效")
            # 在 Maya 中显示错误信息
            om.MGlobal.displayError("顶点距离计算错误")
            return 5

        # 分配权重（优化循环结构）
        logger.info("开始分配权重...")
        # 遍历目标网格的顶点和对应的距离
        for vtx, distance in zip(mesh.vtx, distances):
            if distance > 0:
                # 根据距离和最大距离计算权重
                weight = 1 / max_distance * (max_distance - distance)
                # 使用 skinPercent 函数为顶点分配权重
                pm.skinPercent(
                    skin,
                    vtx,
                    transformValue=[(joint, weight)],
                    normalize=False,
                    zeroRemainingInfluences=False
                )

        logger.info("权重分配完成")
        # 在 Maya 中显示信息提示权重分配完成
        om.MGlobal.displayInfo("蒙皮权重分配完成")
        return 0

    except Exception as e:
        # 若处理过程中发生未预期错误，记录异常信息
        logger.exception("处理过程中发生未预期错误")
        # 在 Maya 中显示错误信息
        om.MGlobal.displayError(f"错误发生: {str(e)}")
        return 6


if __name__ == "__main__":
    try:
        # 交互式选择验证
        # 获取当前在 Maya 中选中的对象
        selection = pm.selected()
        if len(selection) != 3:
            # 若选中的对象数量不为 3，弹出确认对话框提示用户选择三个对象
            pm.confirmDialog(
                title="错误",
                message="请按顺序选择：\n1. 目标网格\n2. 参考网格\n3. 关节",
                button=["OK"]
            )
            # 抛出 ValueError 异常
            raise ValueError("需要选择三个对象")

        # 执行权重分配
        result = dis_weight(
            mesh=selection[0],
            def_mesh=selection[1],
            joint=selection[2]
        )

        # 显示结果
        if result == 0:
            # 若权重分配成功，弹出确认对话框提示用户
            pm.confirmDialog(
                title="完成",
                message="权重分配成功完成",
                button=["OK"]
            )
        else:
            # 若权重分配失败，弹出确认对话框显示错误代码
            pm.confirmDialog(
                title="错误",
                message=f"操作失败，错误代码: {result}",
                button=["OK"]
            )

    except Exception as e:
        # 若主流程发生错误，记录异常信息
        logger.exception("主流程错误")
        # 在 Maya 中显示致命错误信息
        om.MGlobal.displayError(f"致命错误: {str(e)}")